import os
import xml.dom.minidom as xmldom

import cv2
import numpy as np


def filter_out(src_frame):
    if src_frame is not None:
        hsv = cv2.cvtColor(src_frame, cv2.COLOR_BGR2HSV)
        lower_val = np.array([0, 0, 0])
        upper_val = np.array([179, 255, 127])
        mask = cv2.inRange(hsv, lower_val, upper_val)
        mask_inv = cv2.bitwise_not(mask)
        return mask_inv


if __name__ == '__main__':
    '''
    展示labelImg 的结果
    '''
    # rootPath = "../validatecode_dec/gowhere_validate"

    rootPath = "./labeled"
    files = os.listdir(rootPath)
    windowName = "img"

    sum_map = {}

    for file in files:
        if file.endswith(".xml"):
            xmlfname = os.path.join(rootPath, file)
            fname = os.path.join(rootPath, file.replace(".xml", ".png"))
            domobj = xmldom.parse(xmlfname)
            elementobjs = domobj.documentElement

            subobjs = elementobjs.getElementsByTagName("object")

            img = cv2.imread(fname)
            cv2.threshold(img, 180, 255, cv2.THRESH_BINARY, img)
            img = filter_out(img)

            for subobj in subobjs:
                name = subobj.getElementsByTagName("name")[0].firstChild.data
                bndbox = subobj.getElementsByTagName("bndbox")[0]

                xmin = int(bndbox.getElementsByTagName("xmin")[0].firstChild.data)
                ymin = int(bndbox.getElementsByTagName("ymin")[0].firstChild.data)
                xmax = int(bndbox.getElementsByTagName("xmax")[0].firstChild.data)
                ymax = int(bndbox.getElementsByTagName("ymax")[0].firstChild.data)

                # print(name, xmin, ymin, xmax, ymax)

                if sum_map.__contains__(name):
                    sum_map_count = sum_map[name] + 1
                    sum_map[name] = sum_map_count
                else:
                    sum_map[name] = 1
                cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), thickness=1)

            cv2.imshow(windowName, img)
            cv2.resizeWindow(windowName, 600, 600)
            cv2.waitKey(1000)

    print("统计模板过数----------------")

    print(sum_map)

    cv2.destroyAllWindows()
    exit(0)
